import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.linalg import solve_continuous_are
import logging
import time

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger('DroneSimulation')

# 无人机参数
m = 1.0  # 质量 (kg)
g = 9.81  # 重力加速度 (m/s^2)
Ixx = 0.1  # 转动惯量 (kg*m^2)
Iyy = 0.1
Izz = 0.1


# 状态向量: [x, y, z, phi, theta, psi, vx, vy, vz, p, q, r]
# 控制输入: [总推力, 扭矩x, 扭矩y, 扭矩z]

class Drone:
    def __init__(self):
        # 初始状态 - 明确指定为浮点数
        self.state = np.array([0.0, 0.0, 0.0,  # 位置 (x, y, z)
                               0.0, 0.0, 0.0,  # 欧拉角 (φ, θ, ψ)
                               0.0, 0.0, 0.0,  # 速度 (vx, vy, vz)
                               0.0, 0.0, 0.0], dtype=np.float64)  # 角速度 (p, q, r)

        # 目标状态 (悬停状态) - 明确指定为浮点数
        self.target_state = np.array([0.0, 0.0, 5.0,  # 目标高度5米
                                      0.0, 0.0, 0.0,
                                      0.0, 0.0, 0.0,
                                      0.0, 0.0, 0.0], dtype=np.float64)

        # 计算LQR控制器增益
        self.K = self.compute_lqr_gain()

        # 日志计数器
        self.log_counter = 0
        self.log_interval = 100  # 每100步记录一次日志

    def compute_lqr_gain(self):
        # 系统动力学线性化矩阵 (在悬停点附近)
        A = np.zeros((12, 12), dtype=np.float64)
        A[0:3, 6:9] = np.eye(3)
        A[3:6, 9:12] = np.eye(3)
        A[6, 4] = g  # 由于θ角引起的x加速度
        A[7, 3] = -g  # 由于φ角引起的y加速度
        A[8, 8] = 0  # z加速度由推力控制

        # 控制输入矩阵
        B = np.zeros((12, 4), dtype=np.float64)
        B[8, 0] = 1 / m  # 推力对z加速度的影响
        B[9:12, 1:4] = np.array([[1 / Ixx, 0, 0],
                                 [0, 1 / Iyy, 0],
                                 [0, 0, 1 / Izz]], dtype=np.float64)

        # LQR权重矩阵
        Q = np.diag([10, 10, 100,  # 位置权重
                     10, 10, 10,  # 角度权重
                     1, 1, 10,  # 速度权重
                     1, 1, 1])  # 角速度权重

        R = np.diag([0.1, 1, 1, 1])  # 控制输入权重

        # 求解连续时间代数Riccati方程
        P = solve_continuous_are(A, B, Q, R)

        # 计算最优增益矩阵
        K = np.linalg.inv(R) @ B.T @ P

        return K

    def dynamics(self, state, u):
        try:
            # 无人机动力学模型
            x, y, z, phi, theta, psi, vx, vy, vz, p, q, r = state

            # 控制输入: [总推力, 扭矩x, 扭矩y, 扭矩z]
            F, tau_x, tau_y, tau_z = u

            # 位置导数
            dx = vx
            dy = vy
            dz = vz

            # 避免角度奇异点
            if abs(np.cos(theta)) < 1e-5:
                theta = np.sign(theta) * (np.pi / 2 - 1e-5)

            # 欧拉角导数
            dphi = p + (q * np.sin(phi) + r * np.cos(phi)) * np.tan(theta)
            dtheta = q * np.cos(phi) - r * np.sin(phi)
            dpsi = (q * np.sin(phi) + r * np.cos(phi)) / np.cos(theta)

            # 速度导数
            dvx = -F / m * (np.sin(phi) * np.sin(psi) + np.cos(phi) * np.cos(psi) * np.sin(theta))
            dvy = F / m * (np.cos(psi) * np.sin(phi) - np.cos(phi) * np.sin(psi) * np.sin(theta))
            dvz = F / m * (np.cos(phi) * np.cos(theta)) - g

            # 角速度导数
            dp = (tau_x + (Iyy - Izz) * q * r) / Ixx
            dq = (tau_y + (Izz - Ixx) * p * r) / Iyy
            dr = (tau_z + (Ixx - Iyy) * p * q) / Izz

            return np.array([dx, dy, dz, dphi, dtheta, dpsi, dvx, dvy, dvz, dp, dq, dr], dtype=np.float64)

        except Exception as e:
            logger.error(f"Dynamics calculation error: {e}")
            return np.zeros(12, dtype=np.float64)

    def update(self, dt):
        # 计算控制输入
        error = self.state - self.target_state
        u = -self.K @ error

        # 确保推力非负
        u[0] = max(u[0], 0)

        # 使用RK4积分更新状态
        k1 = self.dynamics(self.state, u)
        k2 = self.dynamics(self.state + 0.5 * dt * k1, u)
        k3 = self.dynamics(self.state + 0.5 * dt * k2, u)
        k4 = self.dynamics(self.state + dt * k3, u)

        self.state += dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6

        # 记录日志
        self.log_counter += 1
        if self.log_counter % self.log_interval == 0:
            pos_error = np.linalg.norm(self.state[0:3] - self.target_state[0:3])
            att_error = np.linalg.norm(self.state[3:6] - self.target_state[3:6])
            logger.info(f"Time: {self.log_counter * dt:.2f}s, "
                        f"Position Error: {pos_error:.4f}m, "
                        f"Attitude Error: {np.degrees(att_error):.4f}°, "
                        f"Thrust: {u[0]:.4f}N")

        return self.state, u


# 创建无人机实例
drone = Drone()

# 模拟参数
dt = 0.01  # 时间步长
total_time = 10  # 总模拟时间
steps = int(total_time / dt)

# 存储历史数据
history = {
    'time': np.zeros(steps, dtype=np.float64),
    'position': np.zeros((steps, 3), dtype=np.float64),
    'attitude': np.zeros((steps, 3), dtype=np.float64),
    'velocity': np.zeros((steps, 3), dtype=np.float64),
    'angular_velocity': np.zeros((steps, 3), dtype=np.float64),
    'control': np.zeros((steps, 4), dtype=np.float64)
}

# 运行模拟
logger.info("Starting simulation...")
start_time = time.time()

for i in range(steps):
    state, control = drone.update(dt)
    history['time'][i] = i * dt
    history['position'][i] = state[0:3]
    history['attitude'][i] = state[3:6]
    history['velocity'][i] = state[6:9]
    history['angular_velocity'][i] = state[9:12]
    history['control'][i] = control

elapsed_time = time.time() - start_time
logger.info(f"Simulation completed in {elapsed_time:.2f} seconds")

# 可视化结果
fig, axes = plt.subplots(4, 2, figsize=(14, 12))

# 位置随时间变化
axes[0, 0].plot(history['time'], history['position'][:, 0], label='X')
axes[0, 0].plot(history['time'], history['position'][:, 1], label='Y')
axes[0, 0].plot(history['time'], history['position'][:, 2], label='Z')
axes[0, 0].set_xlabel('Time (s)')
axes[0, 0].set_ylabel('Position (m)')
axes[0, 0].legend()
axes[0, 0].set_title('Position vs Time')
axes[0, 0].grid(True)

# 姿态随时间变化
axes[0, 1].plot(history['time'], np.degrees(history['attitude'][:, 0]), label='Roll')
axes[0, 1].plot(history['time'], np.degrees(history['attitude'][:, 1]), label='Pitch')
axes[0, 1].plot(history['time'], np.degrees(history['attitude'][:, 2]), label='Yaw')
axes[0, 1].set_xlabel('Time (s)')
axes[0, 1].set_ylabel('Angle (degrees)')
axes[0, 1].legend()
axes[0, 1].set_title('Attitude vs Time')
axes[0, 1].grid(True)

# 速度随时间变化
axes[1, 0].plot(history['time'], history['velocity'][:, 0], label='Vx')
axes[1, 0].plot(history['time'], history['velocity'][:, 1], label='Vy')
axes[1, 0].plot(history['time'], history['velocity'][:, 2], label='Vz')
axes[1, 0].set_xlabel('Time (s)')
axes[1, 0].set_ylabel('Velocity (m/s)')
axes[1, 0].legend()
axes[1, 0].set_title('Velocity vs Time')
axes[1, 0].grid(True)

# 角速度随时间变化
axes[1, 1].plot(history['time'], history['angular_velocity'][:, 0], label='p')
axes[1, 1].plot(history['time'], history['angular_velocity'][:, 1], label='q')
axes[1, 1].plot(history['time'], history['angular_velocity'][:, 2], label='r')
axes[1, 1].set_xlabel('Time (s)')
axes[1, 1].set_ylabel('Angular Velocity (rad/s)')
axes[1, 1].legend()
axes[1, 1].set_title('Angular Velocity vs Time')
axes[1, 1].grid(True)

# 控制输入随时间变化
axes[2, 0].plot(history['time'], history['control'][:, 0], label='Thrust')
axes[2, 0].set_xlabel('Time (s)')
axes[2, 0].set_ylabel('Thrust (N)')
axes[2, 0].legend()
axes[2, 0].set_title('Thrust vs Time')
axes[2, 0].grid(True)

axes[2, 1].plot(history['time'], history['control'][:, 1], label='Torque X')
axes[2, 1].plot(history['time'], history['control'][:, 2], label='Torque Y')
axes[2, 1].plot(history['time'], history['control'][:, 3], label='Torque Z')
axes[2, 1].set_xlabel('Time (s)')
axes[2, 1].set_ylabel('Torque (N·m)')
axes[2, 1].legend()
axes[2, 1].set_title('Torques vs Time')
axes[2, 1].grid(True)

# 误差随时间变化
pos_error = np.linalg.norm(history['position'] - np.array([0, 0, 5]), axis=1)
att_error = np.linalg.norm(history['attitude'], axis=1)

axes[3, 0].plot(history['time'], pos_error, label='Position Error')
axes[3, 0].set_xlabel('Time (s)')
axes[3, 0].set_ylabel('Error (m)')
axes[3, 0].legend()
axes[3, 0].set_title('Position Error vs Time')
axes[3, 0].grid(True)

axes[3, 1].plot(history['time'], np.degrees(att_error), label='Attitude Error')
axes[3, 1].set_xlabel('Time (s)')
axes[3, 1].set_ylabel('Error (degrees)')
axes[3, 1].legend()
axes[3, 1].set_title('Attitude Error vs Time')
axes[3, 1].grid(True)

plt.tight_layout()
plt.show()

# 创建动画
fig_anim = plt.figure(figsize=(10, 8))
ax_anim = fig_anim.add_subplot(111, projection='3d')


def update_anim(frame):
    ax_anim.clear()

    # 绘制轨迹
    ax_anim.plot(history['position'][:frame, 0],
                 history['position'][:frame, 1],
                 history['position'][:frame, 2],
                 'b-', alpha=0.5)

    # 绘制当前无人机位置
    ax_anim.scatter(history['position'][frame, 0],
                    history['position'][frame, 1],
                    history['position'][frame, 2],
                    c='red', marker='o', s=100)

    # 绘制目标位置
    ax_anim.scatter(0, 0, 5, c='green', marker='x', s=100)

    ax_anim.set_xlim([-1, 1])
    ax_anim.set_ylim([-1, 1])
    ax_anim.set_zlim([0, 10])
    ax_anim.set_xlabel('X (m)')
    ax_anim.set_ylabel('Y (m)')
    ax_anim.set_zlabel('Z (m)')
    ax_anim.set_title(f'Drone Hover Control (Time: {history["time"][frame]:.2f}s)')

    return ax_anim


# 创建动画
ani = FuncAnimation(fig_anim, update_anim, frames=steps, interval=50, blit=False)
plt.show()